# 设备硬件信息采集
import subprocess
import json
import re

class DeviceInfoCollector(object):

    def __init__(self):
        self.device_info_map = {}

    @staticmethod
    def run_command(command):
        try:
            # 执行linux命令
            result = subprocess.run(command, capture_output=True, text=True)
            if result.returncode == 0:
                rs_str = result.stdout.strip()
                return rs_str
        except:
            pass

    def get_baseboard_info(self):
        """获取主板信息"""
        baseboard_manufacturer = self.run_command(['sudo', 'dmidecode', '-s', 'baseboard-manufacturer']) # 主板制造商
        baseboard_product_name = self.run_command(['sudo', 'dmidecode', '-s', 'baseboard-product-name']) # 主板产品名称
        baseboard_version = self.run_command(['sudo', 'dmidecode', '-s', 'baseboard-version']) # 主板版本
        baseboard_serial_number = self.run_command(['sudo', 'dmidecode', '-s', 'baseboard-serial-number']) # 主板序列号
        self.device_info_map['baseboard_manufacturer'] = baseboard_manufacturer
        self.device_info_map['baseboard_product_name'] = baseboard_product_name
        self.device_info_map['baseboard_version'] = baseboard_version
        self.device_info_map['baseboard_serial_number'] = baseboard_serial_number

    def get_system_info(self):
        """获取系统信息"""
        system_manufacturer = self.run_command(['sudo', 'dmidecode', '-s', 'system-manufacturer'])  # 系统制造商
        system_product_name = self.run_command(['sudo', 'dmidecode', '-s', 'system-product-name'])  # 系统产品名称
        system_version = self.run_command(['sudo', 'dmidecode', '-s', 'system-version'])  # 系统版本
        system_serial_number = self.run_command(['sudo', 'dmidecode', '-s', 'system-serial-number'])  # 系统序列号
        system_uuid = self.run_command(['sudo', 'dmidecode', '-s', 'system-uuid'])  # 系统UUID

        self.device_info_map['system_manufacturer'] = system_manufacturer
        self.device_info_map['system_product_name'] = system_product_name
        self.device_info_map['system_version'] = system_version
        self.device_info_map['system_serial_number'] = system_serial_number
        self.device_info_map['system_uuid'] = system_uuid

    def get_cpu_info(self):
        """获取CPU信息 - 统一数组格式，使用硬件ID"""
        # 使用dmidecode获取所有处理器信息
        dmidecode_output = self.run_command(['sudo', 'dmidecode', '-t', 'processor'])

        processors_info = []
        if dmidecode_output:
            # 直接使用fallback方法解析，因为您的输出格式是标准的
            processors_info = self.parse_fallback_processor_info(dmidecode_output)

        # 如果还是没有获取到信息，使用最后的fallback
        if not processors_info:
            processors_info = self.get_processor_info_fallback()

        # 只保存处理器数组信息
        self.device_info_map['processors'] = processors_info

    def parse_fallback_processor_info(self, output):
        """解析非标准格式的处理器信息"""
        processors_info = []

        # 按行解析查找处理器相关信息
        lines = output.split('\n')
        current_processor = {}

        for line in lines:
            line = line.strip()

            if 'Processor Information' in line:
                if current_processor:
                    processors_info.append(current_processor)
                current_processor = {'cpu_id': f'proc_{len(processors_info)}'}
            elif line.startswith('Socket Designation:') and current_processor:
                current_processor['socket'] = line.split(':', 1)[1].strip()
            elif line.startswith('Manufacturer:') and current_processor:
                current_processor['manufacturer'] = line.split(':', 1)[1].strip()
            elif line.startswith('Version:') and current_processor:
                current_processor['version'] = line.split(':', 1)[1].strip()
            elif line.startswith('Max Speed:') and current_processor:
                current_processor['max_speed'] = line.split(':', 1)[1].strip()
            elif line.startswith('Current Speed:') and current_processor:
                current_processor['current_speed'] = line.split(':', 1)[1].strip()
            elif line.startswith('Status:') and current_processor:
                status = line.split(':', 1)[1].strip()
                if status != 'Unpopulated':
                    current_processor['status'] = status
                else:
                    current_processor = {}  # 跳过未填充的处理器
            elif line.startswith('Core Count:') and current_processor:
                current_processor['core_count'] = line.split(':', 1)[1].strip()
            elif line.startswith('Thread Count:') and current_processor:
                current_processor['thread_count'] = line.split(':', 1)[1].strip()
            elif line.startswith('ID:') and current_processor:
                current_processor['processor_id'] = line.split(':', 1)[1].strip()
            elif line.startswith('Voltage:') and current_processor:
                current_processor['voltage'] = line.split(':', 1)[1].strip()
            elif line.startswith('External Clock:') and current_processor:
                current_processor['external_clock'] = line.split(':', 1)[1].strip()
            elif line.startswith('Upgrade:') and current_processor:
                current_processor['upgrade'] = line.split(':', 1)[1].strip()
            elif line.startswith('Family:') and current_processor:
                current_processor['family'] = line.split(':', 1)[1].strip()

        if current_processor:
            processors_info.append(current_processor)

        return processors_info

    def get_processor_info_fallback(self):
        """最后的fallback方法 - 处理-s参数返回的重复信息"""
        processor_manufacturer = self.run_command(['sudo', 'dmidecode', '-s', 'processor-manufacturer'])
        processor_version = self.run_command(['sudo', 'dmidecode', '-s', 'processor-version'])
        processor_frequency = self.run_command(['sudo', 'dmidecode', '-s', 'processor-frequency'])

        processors_info = []

        # 处理可能的重复信息
        if processor_manufacturer:
            manufacturers = self.split_repeated_info(processor_manufacturer)
            versions = self.split_repeated_info(processor_version) if processor_version else []
            frequencies = self.split_repeated_info(processor_frequency) if processor_frequency else []

            # 确保所有列表长度一致
            max_len = max(len(manufacturers), len(versions), len(frequencies))

            for i in range(max_len):
                processor_info = {
                    'cpu_id': f'fallback_{i}',
                    'manufacturer': manufacturers[i] if i < len(manufacturers) else 'Unknown',
                    'version': versions[i] if i < len(versions) else 'Unknown',
                    'max_speed': frequencies[i] if i < len(frequencies) else 'Unknown',
                    'socket': f'CPU{i + 1}'
                }
                processors_info.append(processor_info)

        return processors_info if processors_info else [{'cpu_id': 'unknown', 'manufacturer': 'Unknown'}]

    def split_repeated_info(self, info_string):
        """拆分重复的信息字符串"""
        if not info_string:
            return []

        # 先尝试按常见分隔符拆分
        if '\n' in info_string:
            items = [item.strip() for item in info_string.split('\n') if item.strip()]
        else:
            # 如果是重复的相同字符串，尝试识别模式
            # 例如: "Intel(R) Corporation Intel(R) Corporation" -> ["Intel(R) Corporation", "Intel(R) Corporation"]
            words = info_string.split()
            if len(words) % 2 == 0:
                mid = len(words) // 2
                first_half = ' '.join(words[:mid])
                second_half = ' '.join(words[mid:])
                if first_half == second_half:
                    return [first_half, second_half]

            # 如果无法识别模式，返回原字符串作为单个项目
            items = [info_string]

        # 去重但保持顺序和数量（因为相同型号的多个CPU是正常的）
        return items

    def split_processor_blocks(self, output):
        """分割处理器信息块"""
        blocks = []
        current_block = ""
        in_processor_block = False

        for line in output.split('\n'):
            if line.startswith('Handle ') and 'DMI type 4' in line:
                if current_block:
                    blocks.append(current_block)
                current_block = line + '\n'
                in_processor_block = True
            elif line.strip() == 'Processor Information':
                if current_block:
                    current_block += line + '\n'
                in_processor_block = True
            elif in_processor_block:
                current_block += line + '\n'
                # 如果遇到空行且已经有了基本信息，可能是块结束
                if line.strip() == '' and 'Processor Information' in current_block:
                    continue

        if current_block:
            blocks.append(current_block)

        return blocks

    def parse_processor_info(self, block):
        """解析单个处理器信息"""
        info = {}

        # 首先提取Handle ID作为硬件ID
        lines = block.split('\n')
        for line in lines:
            if line.startswith('Handle '):
                # 提取Handle ID，格式如: Handle 0x0004, DMI type 4, 48 bytes
                handle_parts = line.split(',')[0]  # 取Handle部分
                handle_id = handle_parts.split()[1]  # 取ID部分，如0x0004
                info['cpu_id'] = handle_id
                break

        # 解析其他信息
        for line in lines:
            line = line.strip()

            if line.startswith('Manufacturer:'):
                info['manufacturer'] = line.split(':', 1)[1].strip()
            elif line.startswith('Version:'):
                info['version'] = line.split(':', 1)[1].strip()
            elif line.startswith('Max Speed:'):
                info['max_speed'] = line.split(':', 1)[1].strip()
            elif line.startswith('Current Speed:'):
                info['current_speed'] = line.split(':', 1)[1].strip()
            elif line.startswith('Status:'):
                info['status'] = line.split(':', 1)[1].strip()
            elif line.startswith('Socket Designation:'):
                info['socket'] = line.split(':', 1)[1].strip()
            elif line.startswith('Core Count:'):
                info['core_count'] = line.split(':', 1)[1].strip()
            elif line.startswith('Thread Count:'):
                info['thread_count'] = line.split(':', 1)[1].strip()
            elif line.startswith('ID:'):
                info['processor_id'] = line.split(':', 1)[1].strip()

        return info

    def get_memory_info(self):
        """使用dmidecode获取内存信息作为备选方案"""
        memory_devices = []

        # 获取详细内存条信息
        dmidecode_output = self.run_command(['sudo', 'dmidecode', '-t', 'memory'])

        if dmidecode_output:
            lines = dmidecode_output.split('\n')
            current_device = {}
            in_memory_device = False

            for line in lines:
                line = line.strip()

                if 'Memory Device' in line:
                    if current_device and 'size' in current_device:
                        memory_devices.append(current_device)
                    current_device = {}
                    in_memory_device = True
                    continue

                if in_memory_device and ':' in line:
                    key, value = line.split(':', 1)
                    key = key.strip()
                    value = value.strip()

                    if key == 'Size' and 'No Module Installed' not in value:
                        current_device['size'] = value  # 内存容量
                    elif key == 'Speed':
                        current_device['speed'] = value  # 内存频率
                    elif key == 'Manufacturer':
                        current_device['vendor'] = value  # 内存厂商
                    elif key == 'Part Number':
                        current_device['product'] = value  # 产品型号
                    elif key == 'Serial Number':
                        current_device['serial'] = value  # 序列号
                    elif key == 'Locator':
                        current_device['slot'] = value  # 内存插槽位置
                    elif key == 'Type':
                        current_device['type'] = value  # 内存类型(DDR4/DDR5等)
                    elif key == 'Data Width':
                        current_device['width'] = value  # 数据宽度

            # 添加最后一个设备
            if current_device and 'size' in current_device:
                memory_devices.append(current_device)

        self.device_info_map['memory_devices'] = memory_devices

    def get_disk_info(self):
        """获取物理硬盘硬件信息"""
        physical_disks = []

        lshw_output = self.run_command(['sudo', 'lshw', '-class', 'disk'])

        if lshw_output:
            disk_blocks = self.split_lshw_disk_blocks(lshw_output)

            for block in disk_blocks:
                disk_info = self.parse_lshw_disk_output(block)
                if disk_info:
                    physical_disks.append(disk_info)

        self.device_info_map['disk_devices'] = physical_disks

    def split_lshw_disk_blocks(self, output):
        """分割lshw磁盘设备块"""
        blocks = []
        current_block = ""

        for line in output.split('\n'):
            if line.strip().startswith('*-namespace'):
                if current_block:
                    blocks.append(current_block)
                current_block = line + '\n'
            elif current_block:
                current_block += line + '\n'

        # 添加最后一个块
        if current_block:
            blocks.append(current_block)

        return blocks

    def parse_lshw_disk_output(self, block):
        """解析lshw硬盘输出信息"""
        info = {}

        for line in block.split('\n'):
            line = line.strip()

            if ':' not in line:
                continue

            key, value = line.split(':', 1)
            key = key.strip()
            value = value.strip()

            if key == 'description':
                info['description'] = value
            elif key == 'product':
                info['model'] = value
            elif key == 'vendor':
                info['vendor'] = value
            elif key == 'physical id':
                info['physical_id'] = value
            elif key == 'bus info':
                info['bus_info'] = value
            elif key == 'logical name':
                info['device'] = value
            elif key == 'version':
                info['version'] = value
            elif key == 'serial':
                info['serial'] = value
            elif key == 'size':
                info['capacity'] = value
            elif key == 'configuration':
                info['configuration'] = value

        return info if info else None

    def get_network_info(self):
        """获取物理网卡硬件信息"""
        networks = self.get_network_info_lshw()
        self.device_info_map['network_devices'] = networks

    def get_network_info_lshw(self):
        """使用lshw获取物理网卡信息"""
        networks = []

        lshw_output = self.run_command(['sudo', 'lshw', '-class', 'network'])

        if lshw_output:
            network_blocks = self.split_lshw_network_blocks(lshw_output)

            for block in network_blocks:
                # 更精确的过滤：只跳过虚拟设备或无物理ID的设备
                if self.should_skip_network_device(block):
                    continue

                network_info = self.parse_lshw_network_output(block)
                if network_info:
                    networks.append(network_info)

        return networks

    def should_skip_network_device(self, block):
        """判断是否应该跳过该网络设备"""
        lines = block.split('\n')

        for line in lines:
            line = line.strip()
            if ':' in line:
                key, value = line.split(':', 1)
                key = key.strip()
                value = value.strip()
                # 跳过configuration没有ip描述的设备
                if key == 'configuration' and not value:
                    return True
                if key == 'configuration' and 'ip=' not in value:
                    return True
                # 跳过虚拟网络设备
                if key == 'description' and 'virtual' in value.lower():
                    return True
                # 跳过没有物理ID的设备（通常是虚拟设备）
                if key == 'physical id' and not value:
                    return True
                # 跳过没有物理ID的设备（通常是虚拟设备）
                if key == 'bus info' and not value:
                    return True
                # 跳过无线网卡（如果需要的话）
                if key == 'description' and 'wireless' in value.lower():
                    return True

        return False

    def split_lshw_network_blocks(self, output):
        """分割lshw网络设备块"""
        blocks = []
        current_block = ""

        for line in output.split('\n'):
            if line.strip().startswith('*-network'):
                if current_block:
                    blocks.append(current_block)
                current_block = line + '\n'
            elif current_block:
                current_block += line + '\n'

        # 添加最后一个块
        if current_block:
            blocks.append(current_block)

        return blocks

    def parse_lshw_network_output(self, block):
        """解析lshw网络设备输出"""
        info = {}

        for line in block.split('\n'):
            line = line.strip()

            if ':' not in line:
                continue

            key, value = line.split(':', 1)
            key = key.strip()
            value = value.strip()

            if key == 'description':
                info['description'] = value
            elif key == 'product':
                info['product'] = value
            elif key == 'vendor':
                info['vendor'] = value
            elif key == 'physical id':
                info['physical_id'] = value
            elif key == 'bus info':
                info['bus_info'] = value
            elif key == 'logical name':
                info['interface'] = value
            elif key == 'version':
                info['version'] = value
            elif key == 'serial':
                info['mac_address'] = value
            elif key == 'size':
                info['speed'] = value
            elif key == 'capacity':
                info['max_speed'] = value
            elif key == 'capabilities':
                info['capabilities'] = value
            elif key == 'configuration':
                info['configuration'] = value

        # 如果获取到了接口名，则获取IP和网关信息
        if 'interface' in info:
            ip_info = self.get_network_ip_info(info['interface'])
            info.update(ip_info)  # 将IP信息合并到网卡信息中
        return info if info else None

    def get_network_ip_info(self, interface):
        """获取指定网卡的IP地址和默认网关"""
        ip_info = {}

        # 获取IP地址
        ip_cmd = ['ip', '-4', 'addr', 'show', interface]
        ip_output = self.run_command(ip_cmd)

        if ip_output:
            import re
            ip_match = re.search(r'inet\s+(\d+(?:\.\d+){3})', ip_output)
            if ip_match:
                ip_info['ip_address'] = ip_match.group(1)

                # 只有当网卡有IP时才获取网关（避免重复获取）
                gateway_cmd = ['ip', 'route', 'show', 'dev', interface]
                gateway_output = self.run_command(gateway_cmd)

                if gateway_output:
                    gateway_match = re.search(r'via\s+(\d+(?:\.\d+){3})', gateway_output)
                    if gateway_match:
                        ip_info['gateway'] = gateway_match.group(1)

        return ip_info

    def get_bios_info(self):
        """获取BIOS信息"""
        bios_vendor = self.run_command(['sudo', 'dmidecode', '-s', 'bios-vendor'])  # BIOS厂商
        bios_version = self.run_command(['sudo', 'dmidecode', '-s', 'bios-version'])  # BIOS版本
        bios_release_date = self.run_command(['sudo', 'dmidecode', '-s', 'bios-release-date'])  # BIOS发布日期

        self.device_info_map['bios_vendor'] = bios_vendor
        self.device_info_map['bios_version'] = bios_version
        self.device_info_map['bios_release_date'] = bios_release_date

    def get_os_info(self):
        """获取操作系统信息"""
        # 系统版本
        os_release = self.run_command(['cat', '/etc/os-release'])
        kernel_version = self.run_command(['uname', '-r'])  # 内核版本
        architecture = self.run_command(['uname', '-m'])  # 系统架构
        hostname = self.run_command(['hostname'])  # 主机名
        uptime = self.run_command(['uptime', '-p'])  # 运行时间

        # 解析os-release
        os_name = None
        os_version = None
        if os_release:
            for line in os_release.split('\n'):
                if line.startswith('PRETTY_NAME='):
                    os_name = line.split('=')[1].strip('"')
                elif line.startswith('VERSION='):
                    os_version = line.split('=')[1].strip('"')

        self.device_info_map['os_name'] = os_name
        self.device_info_map['os_version'] = os_version
        self.device_info_map['kernel_version'] = kernel_version
        self.device_info_map['architecture'] = architecture
        self.device_info_map['hostname'] = hostname
        self.device_info_map['uptime'] = uptime


    def p_info(self):
        """输出采集的信息"""
        return self.device_info_map

    def collect(self):
        """采集所有信息"""
        print("正在采集硬件信息...")
        self.get_baseboard_info()
        self.get_system_info()
        self.get_cpu_info()
        self.get_memory_info()
        self.get_disk_info()
        self.get_network_info()
        self.get_bios_info()
        self.get_os_info()
        print("信息采集完成！")

def __main__():
    aa = DeviceInfoCollector()
    aa.collect()
    aa.p_info()